import pandas as pd
import numpy as np
import numba
import os

#SER
def rolling_regression_b1(x):
    # print(x.shape)
    x1=x[:,0] #注意这个0对应的是'HS300_return'列系数
    y1=x[:,1] #注意这个1对应的是'指数y_return'列系数
    beta1=(y1*(x1-x1.mean())).sum()/(((x1-x1.mean())*(x1-x1.mean())).sum()+0.000000001) #从线性回归公式的角度计算beta1
    return beta1
    # t=df.rolling(50,method="table").apply(rolling_regression,raw=True, engine="numba")

def rolling_regression_b0(x):
    # print(x.shape)
    x1=x[:,0] #注意这个0对应的是'HS300_return'列系数
    y1=x[:,1] #注意这个1对应的是'指数y_return'列系数
    beta1=(y1*(x1-x1.mean())).sum()/(((x1-x1.mean())*(x1-x1.mean())).sum()+0.000000001) #从线性回归公式的角度计算beta1
    beta0 = y1.mean() - beta1*x1.mean()
    return beta0

def reg_beta1(x, y, n):
    # try:
        df = pd.concat([x ,y], axis=1)
        t = df.rolling(n, method="table", min_periods=3).apply(rolling_regression_b1, raw=True, engine="numba")
        t.replace([np.inf, -np.inf], np.nan, inplace=True)
        t = t.fillna(0)
        return t.iloc[:,0]
    # except:
    #     print('error')

def reg_beta0(x,y, n):
    df = pd.concat([x ,y], axis=1)
    t = df.rolling(n, method="table", min_periods=3).apply(rolling_regression_b0,raw=True, engine="numba")
    t.replace([np.inf, -np.inf], np.nan, inplace=True)
    t = t.fillna(0)
    return t.iloc[:,0]

def up_down(df, n):
    df.rolling(n, min_periods=1).apply(lambda x: (x > x[-1].values).astype(int) )

def ser_corr(x,y,n):
    #window日滚动相关系数
    return x.rolling(n, min_periods=1).corr(y)

def argmax_roll(df, window=10):
    return df.rolling(window, min_periods=1).apply(np.argmax,raw=True, engine="numba") + 1 

def argmin_roll(df, window=10):
    return df.rolling(window, min_periods=1).apply(np.argmin,raw=True, engine="numba") + 1

def df2ser(df):
    return  pd.Series(df.values, index=df.index)


# all
def ma_short_to_long(df, s, l):
    ma_short = df.rolling(s, min_periods=1).mean() 
    ma_long = df.rolling(l, min_periods=1).mean() 
    return ma_short/ma_long

def cur_to_pre(df, n):
    return df/df.shift(n)

def cur_to_pre_log(df,n):
    return np.log(cur_to_pre(df,n))

def ma_cur_to_pre(df, c, p, n):
    ma_cur = df.rolling(c, min_periods=1).mean()
    ma_pre = df.rolling(p, min_periods=1).mean()
    return ma_cur/ma_pre.shift(n)

def ser_rank(df, n):
    return df.rolling(n, min_periods=1).rank(pct=True)

def ser_expand_rank(df):
    return df.expanding(min_periods=1).rank(pct=True)

def sum_roll(df, n):
    return df.rolling(n, min_periods=1).sum()

def min_roll(df, n):
    return df.rolling(n, min_periods=1).min()

def max_roll(df, n):
    return df.rolling(n, min_periods=1).max()

def ma_roll(df, n):
    return df.rolling(n, min_periods=1).mean()

def std_roll(df, n):
    return df.rolling(n, min_periods=1).std()

def ts_ewma(df):
    return df.ewm(ignore_na=True, adjust= True, alpha=0.9).mean() 

def sign_p(df):
    dfn = df>0
    return dfn.astype(np.float32)

def sign_n(df):
    dfn = df<0
    return dfn.astype(np.float32)

def industry_dict():
    root = r'data\stock_data\daily\other\indyustry'
    col = pd.read_pickle(r'data\stock_data\daily\consentrate_price\close.pkl.gzip').columns
    ind_dic = {}
    file_lst = os.listdir(root)
    for i in file_lst:
        file_i = os.path.join(root, i)
        inds_i = pd.read_pickle(file_i)
        value = inds_i['代码'].values.tolist()
        new_vaule = []
        for v in value:
            if (v[:3] not in ['300', '688']) and (v in col):
                new_vaule.append((v))
        # new_vaule = [v for v in value if v[:3] not in ['300', '688']]
        ind_dic[i[:-9]] = new_vaule
    return ind_dic


class ReadConData():
    def read_data(self, name='close', root =r'data\stock_data\daily\consentrate_price'):    
        data = pd.read_pickle(os.path.join(root, f'{name}.pkl.gzip')).astype(np.float32)
        return data

# industry_dict()

